"""
蚁群优化算法求解旅行商问题（TSP）

问题描述：
    给定一系列城市和每对城市之间的距离，寻找访问每个城市一次且仅一次，
    最后返回起点的最短路径。
    
算法原理：
    蚁群优化算法（Ant Colony Optimization, ACO）是一种受到蚂蚁觅食行为启发的
    概率型算法，用于解决组合优化问题。蚂蚁在寻找食物路径时会释放信息素，
    后续的蚂蚁会倾向于沿着信息素浓度较高的路径前进，这形成了一个正反馈机制，
    使得最优路径上的信息素积累得更快，最终引导所有蚂蚁收敛到最优路径。

核心概念：
    - 信息素（Pheromone）：路径上信息素的量，表示路径的优劣。
    - 启发信息（Heuristic Information）：通常是距离的倒数，表示路径的吸引力。
    - 信息素挥发：防止信息素无限积累，使算法能够"忘记"差的路径。
    - 信息素更新：蚂蚁完成路径后，根据路径质量更新信息素。

应用场景：
    蚁群优化算法适用于解决各种离散优化问题，例如：
    - 旅行商问题（TSP）
    - 车辆路径问题（VRP）
    - 作业车间调度
    - 网络路由
    - 机器人路径规划

作者：斯黄
日期：2025年3月4日
"""

import numpy as np
import matplotlib.pyplot as plt
import random

# 城市坐标数据
cities = np.array([
    (116.4074, 39.9042),  # 北京
    (121.4737, 31.2304),  # 上海
    (113.2644, 23.1291),  # 广州
    (114.0579, 22.5431),  # 深圳
    (104.0668, 30.6706)   # 成都
])
NUM_CITIES = len(cities)

# 计算距离矩阵
def calculate_distance_matrix(cities):
    num_cities = len(cities)
    dist_matrix = np.zeros((num_cities, num_cities))
    for i in range(num_cities):
        for j in range(i, num_cities):
            dist = np.linalg.norm(cities[i] - cities[j])
            dist_matrix[i, j] = dist
            dist_matrix[j, i] = dist
    return dist_matrix

# 计算路径总距离
def calculate_path_distance(path, dist_matrix):
    total_distance = 0
    for i in range(len(path)):
        total_distance += dist_matrix[path[i], path[(i + 1) % len(path)]]
    return total_distance

# 蚁群优化主程序
def ant_colony_optimization(cities, num_ants, num_iterations, alpha, beta, rho, Q):
    dist_matrix = calculate_distance_matrix(cities)
    num_cities = len(cities)
    
    # 初始化信息素矩阵
    pheromone_matrix = np.ones((num_cities, num_cities)) * 0.1
    
    best_path = None
    best_distance = float('inf')
    history_distances = []

    for iteration in range(num_iterations):
        all_paths = []
        all_path_distances = []

        for ant in range(num_ants):
            current_path = [random.randint(0, num_cities - 1)] # 随机选择起始城市
            unvisited_cities = set(range(num_cities)) - set(current_path)

            while unvisited_cities:
                current_city = current_path[-1]
                probabilities = []
                
                # 计算选择下一个城市的概率
                for next_city in unvisited_cities:
                    pheromone = pheromone_matrix[current_city, next_city] ** alpha
                    heuristic = (1.0 / (dist_matrix[current_city, next_city] + 1e-10)) ** beta # 避免除零
                    probabilities.append(pheromone * heuristic)
                
                total_prob = sum(probabilities)
                if total_prob == 0:
                    next_city = random.choice(list(unvisited_cities))
                else:
                    probabilities = [p / total_prob for p in probabilities]
                    next_city = random.choices(list(unvisited_cities), weights=probabilities, k=1)[0]
                
                current_path.append(next_city)
                unvisited_cities.remove(next_city)
            
            path_distance = calculate_path_distance(current_path, dist_matrix)
            all_paths.append(current_path)
            all_path_distances.append(path_distance)
            
            # 更新全局最优
            if path_distance < best_distance:
                best_distance = path_distance
                best_path = current_path
        
        # 信息素挥发
        pheromone_matrix *= (1 - rho)
        
        # 信息素更新
        for path, dist in zip(all_paths, all_path_distances):
            for i in range(len(path)):
                city_a = path[i]
                city_b = path[(i + 1) % len(path)]
                pheromone_matrix[city_a, city_b] += Q / dist
                pheromone_matrix[city_b, city_a] += Q / dist # 对称问题
        
        history_distances.append(best_distance)

    return best_path, best_distance, history_distances

# 可视化函数
def visualize_tsp_aco(path, distance, generation_history):
    plt.figure(figsize=(12, 6))

    # 绘制路径图
    plt.subplot(1, 2, 1)
    path_coords = cities[path + [path[0]]] # 闭合路径
    plt.plot(path_coords[:, 0], path_coords[:, 1], 'r-o', label=f'Best Path (Distance: {distance:.2f})')
    plt.scatter(cities[:, 0], cities[:, 1], c='blue', label='Cities')
    for i, (x, y) in enumerate(cities):
        plt.text(x, y, str(i), fontsize=12, ha='right')
    plt.title('TSP Solution by Ant Colony Optimization')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.legend()
    plt.grid(True)

    # 绘制收敛曲线
    plt.subplot(1, 2, 2)
    plt.plot(generation_history, color='blue')
    plt.title('Convergence of Best Distance')
    plt.xlabel('Iteration')
    plt.ylabel('Best Distance')
    plt.grid(True)

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 参数设置
    num_ants = 20
    num_iterations = 200
    alpha = 1.0  # 信息素重要程度因子
    beta = 5.0   # 启发信息重要程度因子
    rho = 0.5    # 信息素挥发系数
    Q = 100.0    # 信息素强度

    best_path, min_distance, history = ant_colony_optimization(cities, num_ants, num_iterations, alpha, beta, rho, Q)

    print(f"最终最短路径: {best_path}")
    print(f"最终最短距离: {min_distance:.2f}")

    visualize_tsp_aco(best_path, min_distance, history) 